package com.pearl.sharding.level.table.demo.algorithm;

import org.apache.shardingsphere.infra.util.exception.ShardingSpherePreconditions;
import org.apache.shardingsphere.infra.util.exception.external.sql.type.generic.UnsupportedSQLOperationException;
import org.apache.shardingsphere.sharding.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.sharding.api.sharding.complex.ComplexKeysShardingValue;
import org.apache.shardingsphere.sharding.exception.algorithm.sharding.MismatchedComplexInlineShardingAlgorithmColumnAndValueSizeException;
import org.apache.shardingsphere.sharding.exception.data.NullShardingValueException;

import java.util.*;
import java.util.stream.Collectors;
import java.util.Map.Entry;

/**
 * @author TD
 * @version 1.0
 * @date 2024/2/28
 */
public class MyComplexKeysShardingAlgorithm implements ComplexKeysShardingAlgorithm<Comparable<?>> {

    private static final String ALLOW_RANGE_QUERY_KEY = "allow-range-query-with-inline-sharding";

    private static final String SHARING_COLUMNS_KEY = "sharding-columns";
    private Collection<String> shardingColumns;

    private boolean allowRangeQuery;

    public void init(Properties props) {
        this.allowRangeQuery = this.isAllowRangeQuery(props);
        this.shardingColumns = this.getShardingColumns(props);
    }

    private boolean isAllowRangeQuery(Properties props) {
        return Boolean.parseBoolean(props.getOrDefault(ALLOW_RANGE_QUERY_KEY, Boolean.FALSE.toString()).toString());
    }

    private Collection<String> getShardingColumns(Properties props) {
        String shardingColumns = props.getProperty(SHARING_COLUMNS_KEY, "");
        return shardingColumns.isEmpty() ? Collections.emptyList() : Arrays.asList(shardingColumns.split(","));
    }


    @Override
    public Collection<String> doSharding(Collection<String> availableTargetNames, ComplexKeysShardingValue<Comparable<?>> shardingValue) {
        // 1. 不包含分片键，全路由
        if (!shardingValue.getColumnNameAndRangeValuesMap().isEmpty()) {
            ShardingSpherePreconditions.checkState(this.allowRangeQuery, () -> new UnsupportedSQLOperationException(String.format("Since the property of `%s` is false, inline sharding algorithm can not tackle with range query", "allow-range-query-with-inline-sharding")));
            return availableTargetNames;
        } else {
            // 2. 获取分片键、值
            Map<String, Collection<Comparable<?>>> columnNameAndShardingValuesMap = shardingValue.getColumnNameAndShardingValuesMap();
            // 3. 检查分片键是否为空
            ShardingSpherePreconditions.checkState(this.shardingColumns.isEmpty() || this.shardingColumns.size() == columnNameAndShardingValuesMap.size(),
                    () -> new MismatchedComplexInlineShardingAlgorithmColumnAndValueSizeException(this.shardingColumns.size(), columnNameAndShardingValuesMap.size()));
            // 4. 分片
            return flatten(columnNameAndShardingValuesMap).stream().map(this::doSharding).collect(Collectors.toList());
        }
    }

    // 分片逻辑
    private String doSharding(final Map<String, Comparable<?>> columnNameAndShardingValueMap) {
        // 循环
        for (Entry<String, Comparable<?>> entry : columnNameAndShardingValueMap.entrySet()) {
            String key = entry.getKey();
            Comparable<?> value = entry.getValue();
            ShardingSpherePreconditions.checkNotNull(value, NullShardingValueException::new);
            // 自定义分片计算逻辑
        }
        return "ds_0.t_order_1_1";
    }

    private Collection<Map<String, Comparable<?>>> flatten(final Map<String, Collection<Comparable<?>>> columnNameAndShardingValuesMap) {
        Collection<Map<String, Comparable<?>>> result = new LinkedList<>();
        for (Entry<String, Collection<Comparable<?>>> entry : columnNameAndShardingValuesMap.entrySet()) {
            if (result.isEmpty()) {
                for (Comparable<?> value : entry.getValue()) {
                    Map<String, Comparable<?>> item = new HashMap<>();
                    item.put(entry.getKey(), value);
                    result.add(item);
                }
            } else {
                result = flatten(result, entry.getKey(), entry.getValue());
            }
        }
        return result;
    }

    private Collection<Map<String, Comparable<?>>> flatten(final Collection<Map<String, Comparable<?>>> columnNameAndShardingValueMaps,
                                                           final String columnName, final Collection<Comparable<?>> shardingValues) {
        Collection<Map<String, Comparable<?>>> result = new LinkedList<>();
        for (Map<String, Comparable<?>> each : columnNameAndShardingValueMaps) {
            for (Comparable<?> value : shardingValues) {
                Map<String, Comparable<?>> item = new HashMap<>();
                item.put(columnName, value);
                item.putAll(each);
                result.add(item);
            }
        }
        return result;
    }
}
